package com.lanyun.aicorrectassistant;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import dev.langchain4j.community.model.dashscope.QwenChatModel;
import dev.langchain4j.model.chat.ChatLanguageModel;
import org.apache.poi.xwpf.extractor.XWPFWordExtractor;
import org.apache.poi.xwpf.usermodel.XWPFDocument;

import java.io.FileInputStream;
import java.nio.file.Path;
import java.time.LocalDate;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

public class HomeworkGrader {

    public JSONObject gradeHomework(Path filePath, String userPrompt) throws Exception {
        // 读取文档内容
        String content = readDocxFile(filePath);

        // 自动生成AI指令
        String aiPrompt = buildAutoPrompt(content, userPrompt);

        // 调用AI模型
        String aiResponse = getAIResponse(aiPrompt);

        // 结构化处理结果
        return buildStructuredResult(aiResponse, content);
    }

    private String readDocxFile(Path path) throws Exception {
        try (FileInputStream fis = new FileInputStream(path.toFile());
             XWPFDocument doc = new XWPFDocument(fis);
             XWPFWordExtractor extractor = new XWPFWordExtractor(doc)) {
            return extractor.getText();
        }
    }

    private String buildAutoPrompt(String content, String userPrompt) {
        return String.format(
                "请根据以下作业内容进行智能批改：\n%s\n\n" +
                        "批改要求：\n" +
                        "1. 自动识别作业类型（选择题/策划案/其他）\n" +
                        "2. 统计总分和正确率（如适用）\n" +
                        "3. 结构化分析错误项，包含：题号、错误答案、正确答案、知识点\n" +
                        "4. 给出知识薄弱点总结\n" +
                        "5. 补充用户特别要求：%s\n" +
                        "请用以下标记分隔不同部分：" +
                        "【总分】...【正确率】...【错题分析】...【知识总结】",
                content, userPrompt
        );
    }

    private String getAIResponse(String prompt) {
        ChatLanguageModel qwenModel = QwenChatModel.builder()
                .apiKey("sk-8cc2155ad90049b4bf0ec4bb0ec5cf10")
                .modelName("qwen-max")
                .build();
        return qwenModel.chat(prompt);
    }

    private JSONObject buildStructuredResult(String aiResponse, String rawContent) {
        JSONObject result = new JSONObject();

        // 元数据
        Map<String, Object> metadata = new HashMap<>();
        metadata.put("批改日期", LocalDate.now().toString());
        metadata.put("原始内容摘要", rawContent.substring(0, Math.min(100, rawContent.length())));

        // 核心数据
        Map<String, Object> data = new HashMap<>();
        data.put("总分", extractValue(aiResponse, "【总分】(\\d+)"));
        data.put("正确率", extractValue(aiResponse, "【正确率】(.*?)%"));
        data.put("错题分析", extractErrorAnalysis(aiResponse));
        data.put("知识总结", extractKnowledgeSummary(aiResponse));

        result.put("metadata", metadata);
        result.put("data", data);
        result.put("originalResponse", aiResponse); // 保留原始响应

        return result;
    }

    // 正则提取工具方法
    private String extractValue(String text, String regex) {
        Matcher matcher = Pattern.compile(regex).matcher(text);
        return matcher.find() ? matcher.group(1) : "N/A";
    }

    // 结构化错题分析
    private List<Map<String, String>> extractErrorAnalysis(String text) {
        List<Map<String, String>> errors = new ArrayList<>();
        Pattern pattern = Pattern.compile(
                "题号：(\\d+).*?错误答案：([A-D]).*?正确答案：([A-D]).*?知识点：(.*?)(?=\\n|$)"
        );

        Matcher matcher = pattern.matcher(text);
        while (matcher.find()) {
            Map<String, String> error = new HashMap<>();
            error.put("questionId", matcher.group(1));
            error.put("studentAnswer", matcher.group(2));
            error.put("correctAnswer", matcher.group(3));
            error.put("knowledgePoint", matcher.group(4));
            errors.add(error);
        }
        return errors;
    }

    // 知识总结提取
    private Map<String, Object> extractKnowledgeSummary(String text) {
        Map<String, Object> summary = new HashMap<>();
        Pattern pattern = Pattern.compile("【知识总结】([\\s\\S]*?)(?=【|$)");
        Matcher matcher = pattern.matcher(text);
        if (matcher.find()) {
            String[] points = matcher.group(1).split("\\n");
            summary.put("weakAreas", points);
        }
        return summary;
    }

    public static void main(String[] args) throws Exception {
        HomeworkGrader grader = new HomeworkGrader();
        JSONObject result = grader.gradeHomework(
                Path.of("C:/Users/Jqcc/Desktop/作业.docx"),
                "请重点检查数据准确性"
        );
        System.out.println(JSON.toJSONString(result, true));
    }
}